import os
from multiprocessing import Pool,cpu_count
from typing import List
from PIL import Image

'''
数据集的构建：
1、原数据集：
    cat_and_dog
        Cat
            1.jpg
            2.jpg
            ...
        Dog
            1.jpg
            2.jpg
            ...
2、处理成
    data
        test
            Cat
                1.jpg
                2.jpg
                ...
            Dog
                1.jpg
                2.jpg
                ...
        train
            Cat
                1.jpg
                2.jpg
                ...
            Dog
                1.jpg
                2.jpg
                ...            
'''
class MakeDataSet:
    # 下载的原始数据集存放的路径
    source_path = '../a_CatAndDogSourceData'
    # 训练和测试数据集存放的根目录
    save_path = '../a_CatAndDog'
    # 图片的格式
    image_suffix_list = ('.jpg','.jpeg','.png','.bmp')
    # 训练集和测试集的占比
    train_proportion = 0.8
    test_proportion = 0.2
    # 训练/测试图片所在目录的上级目录（在save_path后的目录）
    train_dir_name = 'train'
    test_dir_name = 'test'

    # 删除原始数据集中的损坏的图片文件
    @staticmethod
    def delete_error_image(image_path):
        try:
            if os.path.exists(image_path):
                os.remove(image_path)
                print(f'删除文件 {image_path} 成功')
            else:
                print(f'{image_path} 文件不存在')
        except Exception as e:
            print(f'其他错误：{e}')

    # 创建目录
    @staticmethod
    def create_dir(dir_path:str):
        if os.path.exists(dir_path):
            return
        try:
            os.makedirs(dir_path)
        except Exception as e:
            raise NotADirectoryError(f'创建此目录失败{dir_path},失败原因：{e}')

    # 获取处理原始数据集的进程数大小
    @classmethod
    def get_pool_size(cls):
        class_name_list = os.listdir(cls.source_path)
        cpu_nums = cpu_count()
        class_nums = len(class_name_list)
        if class_nums*2 <= cpu_nums:
            return class_nums*2
        else:
            return cpu_nums

    ## 进程池回调函数
    @classmethod
    def result_callback(self, result):
        ## 回调函数
        print(f"Task result: {result}")

    ## 进程任务
    @classmethod
    def work_task_process(cls,file_path_list:List[str],data_type:str,data_class:str):
        '''
        功能：保存同一个种类图片的训练数据集和测试数据集
        :param file_path_list: [数据集原始路径图片的路径,]，这里是cat_and_dog/Cat/1.jpg
        :param data_type: test/train 测试数据/训练数据
        :param data_class: 数据的分类，当前是猫狗分类，cat/dog
        :return:
        '''
        ## 创建路径
        dir_path = os.path.join(cls.save_path,data_type,data_class)
        cls.create_dir(dir_path)
        for file_path in file_path_list:
            try:
                Image.open(file_path)
            except Exception:
                print(f'{file_path}图片损坏,不保存此图片')
                cls.delete_error_image(file_path)
                continue
            file_name = os.path.basename(file_path)
            with open(file_path,'rb') as fp1:
                save_file_path = os.path.join(dir_path,file_name)
                with open(save_file_path,'wb') as fp2:
                    fp2.write(fp1.read())
        return f'{data_type}/{data_class}保存了{len(file_path_list)}张图片'

    @classmethod
    def main(cls):
        # 进程池
        pool = Pool(processes= cls.get_pool_size())
        # 原始的数据集：拆封成训练数据集和测试数据集
        class_name_list = os.listdir(cls.source_path)
        for class_name in class_name_list:
            ## 当前分类的所在的路径
            class_path = os.path.join(cls.source_path,class_name)
            # 当前分类下的所有数据
            filename_list = [
                os.path.join(cls.source_path,class_name,name) for name in os.listdir(class_path)
                if os.path.splitext(name)[-1] in cls.image_suffix_list
            ]
            # 训练数据集的占比
            train_num = int(len(filename_list)*cls.train_proportion)
            # 训练数据集
            train_file_list = filename_list[:train_num]
            # 测试数据集
            test_file_list = filename_list[train_num:]
            # 起一个进程去处理训练图片
            train_args = (train_file_list,cls.train_dir_name,class_name)
            # 起一个进程去处理测试图片
            test_args = (test_file_list, cls.test_dir_name, class_name)
            # 动态生成任务
            pool.apply_async(cls.work_task_process, args=train_args, callback=cls.result_callback)
            pool.apply_async(cls.work_task_process, args=test_args, callback=cls.result_callback)
        # 关闭进程池，等待所有任务完成
        pool.close()
        pool.join()

if __name__ == '__main__':
    ## 切分训练数据集和测试数据集
    MakeDataSet.main()
    # print(MakeDataSet.get_classes())